{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tce3stUlHN0L"
      },
      "source": [
        "##### Copyright 2018 The TensorFlow Authors.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MfBg1C5NB3X0"
      },
      "source": [
        "# TPU 사용하기\n",
        "\n",
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td><a target=\"_blank\" href=\"https://www.tensorflow.org/guide/tpu\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\">TensorFlow.org에서 보기</a></td>\n",
        "  <td><a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/ko/guide/tpu.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\">Google Colab에서 실행</a></td>\n",
        "  <td><a target=\"_blank\" href=\"https://github.com/tensorflow/docs-l10n/blob/master/site/ko/guide/tpu.ipynb\">     <img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\">    GitHub에서 소스 보기</a></td>\n",
        "  <td><a href=\"https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/ko/guide/tpu.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\">노트북 다운로드</a></td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ys81cOhXOWUP"
      },
      "source": [
        "Cloud TPU에 대한 실험적인 지원은 현재 Keras 및 Google Colab에 제공됩니다. Colab 노트북을 실행하기 전에 노트북 설정(Runtime &gt; Change runtime type &gt; Hardware accelerator &gt; TPU)을 확인하여 하드웨어 가속기가 TPU인지 확인하세요."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ek5Hop74NVKm"
      },
      "source": [
        "## 설정"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Cw0WRaChRxTL"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "import os\n",
        "import tensorflow_datasets as tfds"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yDWaRxSpwBN1"
      },
      "source": [
        "## TPU 초기화\n",
        "\n",
        "TPU는 일반적으로 사용자 파이썬 프로그램을 실행하는 로컬 프로세스와는 다른 Cloud TPU 작업자(worker)에 있습니다. 따라서 원격 클러스터에 연결하고 TPU를 초기화하려면 일부 초기화 작업을 수행해야 합니다. `TPUClusterResolver`에 대한 `tpu` 인수는 Colab 전용의 특수 주소입니다. Google Compute Engine(GCE)에서 실행 중인 경우, CloudTPU 이름으로 전달해야 합니다."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dCqWMqvtwOLs"
      },
      "source": [
        "참고: TPU 초기화 코드는 프로그램의 시작 부분에 있어야 합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dKPqF8d1wJCV"
      },
      "outputs": [],
      "source": [
        "resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')\n",
        "tf.config.experimental_connect_to_cluster(resolver)\n",
        "# This is the TPU initialization code that has to be at the beginning.\n",
        "tf.tpu.experimental.initialize_tpu_system(resolver)\n",
        "print(\"All devices: \", tf.config.list_logical_devices('TPU'))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mv7kehTZ1Lq_"
      },
      "source": [
        "## 수동 기기 배치\n",
        "\n",
        "TPU가 초기화된 후 수동 기기 배치를 사용하여 단일 TPU 기기에 계산을 배치할 수 있습니다.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XRZ4kMoxBNND"
      },
      "outputs": [],
      "source": [
        "a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])\n",
        "b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])\n",
        "with tf.device('/TPU:0'):\n",
        "  c = tf.matmul(a, b)\n",
        "print(\"c device: \", c.device)\n",
        "print(c)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_NJm-kgFO0cC"
      },
      "source": [
        "## 배포 전략\n",
        "\n",
        "대부분의 경우, 사용자는 데이터를 병렬로 여러 TPU에서 모델을 실행하려고 합니다. 배포 전략은 CPU, GPU 또는 TPU에서 모델을 구동하는 데 사용할 수 있는 추상화입니다. 배포 전략을 바꾸면 모델이 지정된 기기에서 실행됩니다. 자세한 내용은 [배포 전략 가이드](./distributed_training.ipynb)를 참조하세요."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DcDPMZs-9uLJ"
      },
      "source": [
        "먼저, `TPUStrategy` 객체를 만듭니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7SO23K8oRpjI"
      },
      "outputs": [],
      "source": [
        "strategy = tf.distribute.experimental.TPUStrategy(resolver)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JlaAmswWPsU6"
      },
      "source": [
        "계산이 모든 TPU 코어에서 실행될 수 있도록 계산을 복제하기 위해 간단하게 `strategy.run` API에 전달할 수 있습니다. 아래는 모든 코어가 같은 입력 `(a, b)`을 가져와 각 코어에서 독립적으로 matmul을 수행하는 예제입니다. 출력은 모든 복제본의 값이 됩니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-90CL5uFPTOa"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def matmul_fn(x, y):\n",
        "  z = tf.matmul(x, y)\n",
        "  return z\n",
        "\n",
        "z = strategy.run(matmul_fn, args=(a, b))\n",
        "print(z)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uxgYl6kGHJLc"
      },
      "source": [
        "## TPU 기반 분류\n",
        "\n",
        "기본 개념을 배웠으므로 더 구체적인 예를 살펴볼 차례입니다. 이 가이드는 배포 전략 `tf.distribute.experimental.TPUStrategy`을 사용하여 Cloud TPU를 구동하고 Keras 모델을 훈련하는 방법을 보여줍니다.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gKRALGgt_kCo"
      },
      "source": [
        "### Keras 모델 정의하기\n",
        "\n",
        "아래는 Keras를 사용한 MNIST 모델의 정의이며, CPU 또는 GPU에서 사용하는 모델이 변경되지 않았습니다. Keras 모델 작성은 `strategy.scope` 내에 있어야 하므로 각 TPU 기기에서 변수를 작성할 수 있습니다. 코드의 다른 부분은 전략 범위 내에 있을 필요는 없습니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DiBiN-Z_R7P7"
      },
      "outputs": [],
      "source": [
        "def create_model():\n",
        "  return tf.keras.Sequential(\n",
        "      [tf.keras.layers.Conv2D(256, 3, activation='relu', input_shape=(28, 28, 1)),\n",
        "       tf.keras.layers.Conv2D(256, 3, activation='relu'),\n",
        "       tf.keras.layers.Flatten(),\n",
        "       tf.keras.layers.Dense(256, activation='relu'),\n",
        "       tf.keras.layers.Dense(128, activation='relu'),\n",
        "       tf.keras.layers.Dense(10)])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qYOYjYTg_31l"
      },
      "source": [
        "### 입력 데이터세트\n",
        "\n",
        "Cloud TPU를 사용하는 경우, `tf.data.Dataset` API를 효율적으로 사용하는 것이 중요합니다. 데이터를 충분히 빠르게 공급할 수 없다면 Cloud TPU를 사용할 수 없습니다. 데이터세트 성능에 대한 자세한 내용은 [입력 파이프라인 성능 가이드](./data_performance.ipynb)를 참조하세요.\n",
        "\n",
        "가장 간단한 실험(`tf.data.Dataset.from_tensor_slices` 또는 기타 in-graph 데이터 사용)을 제외하고, Dataset에서 읽은 모든 데이터 파일을 Google Cloud Storage(GCS) 버킷에 저장해야 합니다.\n",
        "\n",
        "사용 사례 대부분의 경우, 데이터를 `TFRecord` 형식으로 변환하고 `tf.data.TFRecordDataset`을 사용하여 데이터를 읽는 것이 좋습니다. 이를 수행하는 방법에 대한 자세한 내용은 [TFRecord 및 tf.Example 튜토리얼](../tutorials/load_data/tfrecord.ipynb)을 참조하세요. 엄격한 요구 사항은 아니며, 원하는 경우 다른 데이터세트 리더 (`FixedLengthRecordDataset` 또는`TextLineDataset`)를 사용할 수 있습니다.\n",
        "\n",
        "작은 데이터세트는 `tf.data.Dataset.cache`를 사용하여 메모리에 완전히 로드할 수 있습니다.\n",
        "\n",
        "사용된 데이터 형식과 관계없이 100MB 정도의 큰 파일을 사용하는 것이 좋습니다. 파일을 여는 오버헤드가 상당히 높기 때문에 네트워크 설정에서 특히 중요합니다.\n",
        "\n",
        "이때 `tensorflow_datasets` 모듈을 사용하여 MNIST 훈련 데이터의 사본을 가져와야 합니다. 공개 GCS 버킷에서 사용 가능한 사본을 사용하도록 `try_gcs`이 지정되었습니다. 이를 지정하지 않으면 TPU가 다운로드된 데이터에 액세스할 수 없습니다. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "noAd416KSCo7"
      },
      "outputs": [],
      "source": [
        "def get_dataset(batch_size, is_training=True):\n",
        "  split = 'train' if is_training else 'test'\n",
        "  dataset, info = tfds.load(name='mnist', split=split, with_info=True,\n",
        "                            as_supervised=True, try_gcs=True)\n",
        "\n",
        "  def scale(image, label):\n",
        "    image = tf.cast(image, tf.float32)\n",
        "    image /= 255.0\n",
        "\n",
        "    return image, label\n",
        "\n",
        "  dataset = dataset.map(scale)\n",
        "\n",
        "  # Only shuffle and repeat the dataset in training. The advantage to have a\n",
        "  # infinite dataset for training is to avoid the potential last partial batch\n",
        "  # in each epoch, so users don't need to think about scaling the gradients\n",
        "  # based on the actual batch size.\n",
        "  if is_training:\n",
        "    dataset = dataset.shuffle(10000)\n",
        "    dataset = dataset.repeat()\n",
        "\n",
        "  dataset = dataset.batch(batch_size)\n",
        "\n",
        "  return dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mgUC6A-zCMEr"
      },
      "source": [
        "### Keras 고수준 API를 사용하여 모델 훈련하기\n",
        "\n",
        "Keras fit/compile API를 사용하여 모델을 간단히 훈련할 수 있습니다. TPU에만 해당되는 것은 없으며, 여러 GPU가 있고 `TPUStrategy` 대신 `MirroredStrategy` 사용하는 경우 아래의 같은 코드를 작성합니다. 자세한 내용은 [Keras를 사용한 분산 훈련](https://www.tensorflow.org/tutorials/distribute/keras) 튜토리얼을 확인하세요."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ubmDchPqSIx0"
      },
      "outputs": [],
      "source": [
        "with strategy.scope():\n",
        "  model = create_model()\n",
        "  model.compile(optimizer='adam',\n",
        "                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "                metrics=['sparse_categorical_accuracy'])\n",
        "\n",
        "batch_size = 200\n",
        "steps_per_epoch = 60000 // batch_size\n",
        "validation_steps = 10000 // batch_size\n",
        "\n",
        "train_dataset = get_dataset(batch_size, is_training=True)\n",
        "test_dataset = get_dataset(batch_size, is_training=False)\n",
        "\n",
        "model.fit(train_dataset,\n",
        "          epochs=5,\n",
        "          steps_per_epoch=steps_per_epoch,\n",
        "          validation_data=test_dataset, \n",
        "          validation_steps=validation_steps)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8hSGBIYtUugJ"
      },
      "source": [
        "To reduce python overhead, and maximize the performance of your TPU, try out the **experimental** `experimental_steps_per_execution` argument to `Model.compile`. Here it increases throughput by about 50%:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M6e3aVVLUorL"
      },
      "outputs": [],
      "source": [
        "with strategy.scope():\n",
        "  model = create_model()\n",
        "  model.compile(optimizer='adam',\n",
        "                # Anything between 2 and `steps_per_epoch` could help here.\n",
        "                experimental_steps_per_execution = 50,\n",
        "                loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "                metrics=['sparse_categorical_accuracy'])\n",
        "\n",
        "model.fit(train_dataset,\n",
        "          epochs=5,\n",
        "          steps_per_epoch=steps_per_epoch,\n",
        "          validation_data=test_dataset,\n",
        "          validation_steps=validation_steps)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0rRALBZNCO4A"
      },
      "source": [
        "### 사용자 정의 훈련 루프를 사용하여 모델 훈련하기\n",
        "\n",
        "`tf.function` 및 `tf.distribute` API를 직접 사용하여 모델을 작성하고 훈련할 수도 있습니다. `strategy.experimental_distribute_datasets_from_function` API는 데이터세트 함수가 지정된 데이터세트를 분배하는 데 사용됩니다. 이 경우 데이터세트에 전달된 배치 크기는 전역 배치 크기가 아닌 복제본 배치 크기입니다. 자세한 내용은 [tf.distribute.Strategy를 사용한 사용자 정의 훈련](https://www.tensorflow.org/tutorials/distribute/custom_training) 튜토리얼을 확인하세요.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DxdgXPAL6iFE"
      },
      "source": [
        "먼저, 모델, 데이터세트 및 tf.functions를 작성합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9aHhqwao2Fxi"
      },
      "outputs": [],
      "source": [
        "# Create the model, optimizer and metrics inside strategy scope, so that the\n",
        "# variables can be mirrored on each device.\n",
        "with strategy.scope():\n",
        "  model = create_model()\n",
        "  optimizer = tf.keras.optimizers.Adam()\n",
        "  training_loss = tf.keras.metrics.Mean('training_loss', dtype=tf.float32)\n",
        "  training_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(\n",
        "      'training_accuracy', dtype=tf.float32)\n",
        "\n",
        "# Calculate per replica batch size, and distribute the datasets on each TPU\n",
        "# worker.\n",
        "per_replica_batch_size = batch_size // strategy.num_replicas_in_sync\n",
        "\n",
        "train_dataset = strategy.experimental_distribute_datasets_from_function(\n",
        "    lambda _: get_dataset(per_replica_batch_size, is_training=True))\n",
        "\n",
        "@tf.function\n",
        "def train_step(iterator):\n",
        "  \"\"\"The step function for one training step\"\"\"\n",
        "\n",
        "  def step_fn(inputs):\n",
        "    \"\"\"The computation to run on each TPU device.\"\"\"\n",
        "    images, labels = inputs\n",
        "    with tf.GradientTape() as tape:\n",
        "      logits = model(images, training=True)\n",
        "      loss = tf.keras.losses.sparse_categorical_crossentropy(\n",
        "          labels, logits, from_logits=True)\n",
        "      loss = tf.nn.compute_average_loss(loss, global_batch_size=batch_size)\n",
        "    grads = tape.gradient(loss, model.trainable_variables)\n",
        "    optimizer.apply_gradients(list(zip(grads, model.trainable_variables)))\n",
        "    training_loss.update_state(loss * strategy.num_replicas_in_sync)\n",
        "    training_accuracy.update_state(labels, logits)\n",
        "\n",
        "  strategy.run(step_fn, args=(next(iterator),))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ibi7Z97V6xsQ"
      },
      "source": [
        "그런 다음 훈련 루프를 실행합니다."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1du5cXWt6Vtw"
      },
      "outputs": [],
      "source": [
        "steps_per_eval = 10000 // batch_size\n",
        "\n",
        "train_iterator = iter(train_dataset)\n",
        "for epoch in range(5):\n",
        "  print('Epoch: {}/5'.format(epoch))\n",
        "\n",
        "  for step in range(steps_per_epoch):\n",
        "    train_step(train_iterator)\n",
        "  print('Current step: {}, training loss: {}, accuracy: {}%'.format(\n",
        "      optimizer.iterations.numpy(),\n",
        "      round(float(training_loss.result()), 4),\n",
        "      round(float(training_accuracy.result()) * 100, 2)))\n",
        "  training_loss.reset_states()\n",
        "  training_accuracy.reset_states()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TnZJUM3qIjKu"
      },
      "source": [
        "### `tf.function` 내에서 여러 단계로 성능 향상하기\n",
        "\n",
        "`tf.function` 내에서 여러 단계를 실행하여 성능을 향상할 수 있습니다. `strategy.run` 호출을 `tf.function` 내에 `tf.range`로 래핑하고, AutoGraph는 TPU 작업자에서 `tf.while_loop`로 변환합니다.\n",
        "\n",
        "더 나은 성능을 제공하지만 `tf.function` 내의 단일 단계와 비교하여 보완해야 할 점이 있습니다. `tf.function`에서 여러 단계를 실행하면 유연성이 떨어지므로 단계 내에서 즉시 또는 임의의 파이썬 코드를 실행할 수 없습니다.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2grYvXLzJYkP"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def train_multiple_steps(iterator, steps):\n",
        "  \"\"\"The step function for one training step\"\"\"\n",
        "\n",
        "  def step_fn(inputs):\n",
        "    \"\"\"The computation to run on each TPU device.\"\"\"\n",
        "    images, labels = inputs\n",
        "    with tf.GradientTape() as tape:\n",
        "      logits = model(images, training=True)\n",
        "      loss = tf.keras.losses.sparse_categorical_crossentropy(\n",
        "          labels, logits, from_logits=True)\n",
        "      loss = tf.nn.compute_average_loss(loss, global_batch_size=batch_size)\n",
        "    grads = tape.gradient(loss, model.trainable_variables)\n",
        "    optimizer.apply_gradients(list(zip(grads, model.trainable_variables)))\n",
        "    training_loss.update_state(loss * strategy.num_replicas_in_sync)\n",
        "    training_accuracy.update_state(labels, logits)\n",
        "\n",
        "  for _ in tf.range(steps):\n",
        "    strategy.run(step_fn, args=(next(iterator),))\n",
        "\n",
        "# Convert `steps_per_epoch` to `tf.Tensor` so the `tf.function` won't get \n",
        "# retraced if the value changes.\n",
        "train_multiple_steps(train_iterator, tf.convert_to_tensor(steps_per_epoch))\n",
        "\n",
        "print('Current step: {}, training loss: {}, accuracy: {}%'.format(\n",
        "      optimizer.iterations.numpy(),\n",
        "      round(float(training_loss.result()), 4),\n",
        "      round(float(training_accuracy.result()) * 100, 2)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WBKVhMvWjibf"
      },
      "source": [
        "## 다음 단계\n",
        "\n",
        "- [Google Cloud TPU 설명서](https://cloud.google.com/tpu/docs/) - Google Cloud TPU 설정 및 실행\n",
        "- [TensorFlow를 사용한 분산 훈련](./distributed_training.ipynb) - 배포 전략을 사용하는 방법과 모범 사례를 보여주는 여러 예제에 대한 링크\n",
        "- [TensorFlow 공식 모델](https://github.com/tensorflow/models/tree/master/official) - Cloud TPU와 호환되는 최신 TensorFlow 2.x 모델의 예\n",
        "- [Google Cloud TPU 성능 가이드](https://cloud.google.com/tpu/docs/performance-guide) - 애플리케이션에 대한 Cloud TPU 구성 매개변수를 조정하여 Cloud TPU 성능 향상하기\n",
        "- [Google Cloud TPU 성능 가이드](https://cloud.google.com/tpu/docs/performance-guide) - 애플리케이션에 대한 Cloud TPU 구성 매개변수를 조정하여 Cloud TPU 성능 향상하기."
      ]
    }
  ],
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "collapsed_sections": [],
      "name": "tpu.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
